package partition;

import jdd.util.Array;
import network.Device;
import network.Network;
import network.Rule;
import networkLoader.FBNetwork;
import networkLoader.FabricNetwork;
import org.jgrapht.Graph;
import org.jgrapht.graph.DefaultEdge;
import org.jgrapht.graph.DefaultUndirectedGraph;

import java.util.*;

public class Partitioner {

    /**
     * Create partition with number of nPartition of rsws
     *
     * @param network
     * @param nPartition the number of partitions on rsws
     * @return
     */
    public static HashMap<String, Partition> setupPartitions(Network network, int nPartition) {
        HashMap<String, Partition> nameToPartition = new HashMap<>();
        int nRswPerPartition = 112 * 48 / nPartition;
        int nFswPerPartition = 112 * 4 / nPartition;
        Device[][] partitionDevices = new Device[nPartition][112 * 52 / nPartition];
        for (int iPod = 0; iPod < 112; iPod++) {
            for (int iRsw = 0; iRsw < 48; iRsw++) {
                partitionDevices[(iPod * 48 + iRsw) / nRswPerPartition][(iPod * 48 + iRsw) % nRswPerPartition] = network.getDevice("rsw-" + iPod + "-" + iRsw);
            }
            for (int iFsw = 0; iFsw < 4; iFsw++) {
                partitionDevices[(iPod * 4 + iFsw) / nFswPerPartition][nRswPerPartition + (iPod * 4 + iFsw) % nFswPerPartition] = network.getDevice("fsw-" + iPod + "-" + iFsw);
            }
        }
        for (int iPartition = 0; iPartition < nPartition; iPartition++) {
            Network partition = network.getSubNetwork(Arrays.asList(partitionDevices[iPartition]));
            nameToPartition.put("p-" + iPartition, new Partition("p-" + iPartition, partition));
        }
        for (int i = 0; i < 4; i++) { // spine
            LinkedList<Device> partition = new LinkedList<>();
            for (int j = 0; j < 48; j++) { // spine.ssw
//                partition.add(n.getDevice("default"));
                String ssw = "ssw-" + i + "-" + j;
                partition.add(network.getDevice(ssw));
            }
            Network pt = network.getSubNetwork(partition);
            Partition partition1 = new Partition("spine-" + i, pt);
            nameToPartition.put("spine-" + i, partition1);
        }

        return nameToPartition;
    }

    public static HashMap<String, Partition> setupPartitionWithLTs(Network network) {
        HashMap<String, Partition> nameToPartition = FabricNetwork.setupPartitions(network);
        for (Partition partition : nameToPartition.values()) {
            String ptname = partition.getName();
            if (ptname.startsWith("pod")) {
                String podId = ptname.split("-")[1];
                // inter pod up
                for (int iRsw = 0; iRsw < 48; iRsw++) {
                    for (int iFsw = 0; iFsw < 4; iFsw++) {
                        for (int iPort = 0; iPort < 48; iPort++) {
                            LocalTask localTask = new LocalTask(partition.getNetwork().getDevice("rsw-" + podId + "-" + iRsw),
                                    partition.getNetwork().getDevice("fsw-" + podId + "-" + iFsw).getPort("fsw-" + podId + "-" + iFsw + ">ssw-" + iFsw + "-" + iPort));
                            partition.addLocalTask(localTask);
                        }
                    }
                }
                // inter pod down
                for (int iFsw = 0; iFsw < 4; iFsw++) {
                    for (int iRsw = 0; iRsw < 48; iRsw++) {
                        LocalTask localTaskDown = new LocalTask(
                                partition.getNetwork().getDevice("fsw-" + podId + "-" + iFsw),
                                partition.getNetwork().getDevice("rsw-" + podId + "-" + iRsw).getPort("rsw-" + podId + "-" + iRsw + ">h-" + podId + "-" + iRsw));
                        partition.addLocalTask(localTaskDown);
                    }
                }
                // intra pod
                for (int iRsw = 0; iRsw < 48; iRsw++) {
                    for (int iRsw1 = 0; iRsw1 < 48; iRsw1++) {
                        if (iRsw != iRsw1) {
                            LocalTask localTask = new LocalTask(partition.getNetwork().getDevice("rsw-" + podId + "-" + iRsw),
                                    partition.getNetwork().getDevice("rsw-" + podId + "-" + iRsw1).getPort("rsw-" + podId + "-" + iRsw1 + ">h-" + podId + "-" + iRsw1));
                            partition.addLocalTask(localTask);
                        }
                    }
                }
            }
            if (ptname.startsWith("spine")) {
                for (Device device : partition.getNetwork().getAllDevices()) {
                    String iSpine = ptname.split("-")[1];
                    String iSsw = ptname.split("-")[2];
                    for (int iPod = 0; iPod < 112; iPod++) {
                        if (iPod != Integer.parseInt(iSsw)) {
                            LocalTask localTask = new LocalTask(device,
                                    device.getPort(device.getName() + ">fsw-" + iPod + "-" + iSpine));
                            partition.addLocalTask(localTask);
                        }
                    }
                }
            }

        }
        return nameToPartition;
    }

    public static void main(String[] args) {
        Network n = FBNetwork.getNetwork();
        int[] nPartitions = {1, 1, 2, 4, 7, 8, 14, 16, 28, 56, 112}; // two 1 forbid JIT
//        int[] nPartitions = {112}; // two 1 forbid JIT
        HashMap<Device, HashSet<Rule>> deviceToRules = new HashMap<>();
        for (Rule rule : n.getInitialRules()) {
            if (rule.getMatch().longValue() >> 24 == 1) { // test pod2 as subnet
                deviceToRules.computeIfAbsent(rule.getDevice(), k -> new HashSet<>());
                deviceToRules.get(rule.getDevice()).add(rule);
            }
        }
        for (int nPartition : nPartitions) {
            HashMap<String, Partition> nameToPartition = setupPartitions(n, nPartition);
            for (Partition partition : nameToPartition.values()) {
                ArrayList<Rule> rulesInPartition = new ArrayList<>();
                ArrayList<ArrayList<Rule>> ruless = new ArrayList<>();
                for (Device device : partition.getNetwork().getAllDevices()) {
                    rulesInPartition.addAll(deviceToRules.get(device));
                    ruless.add(new ArrayList<>(deviceToRules.get(device)));
                }
                long s = System.nanoTime();
                partition.runSubnetRuleBatch(ruless);
                long end = (System.nanoTime() - s);
                if (partition.getName().startsWith("p")) {
                    System.out.println("#partition " + nPartition + " " + partition.getName() + " model update takes: " + end);
                }
            }
//            Aggregator aggregator = new Aggregator(nameToPartition);
//            long s = System.nanoTime();
//            aggregator.getNetworkModel();
//            System.out.println("#partition " + nPartition + " aggr gnm takes: " + (System.nanoTime() - s));
        }
    }
}
